;;;;;; code-generator.scm - JVM Code Generator module.    -*- Mode: Scheme -*-
;;;;;; Author: Erik Silkensen <silkense@colorado.edu>
;;;;;; Version: 19 Sep 2009

#lang scheme

(require "parser.scm")
(require "attributes.scm")
(require "class-file.scm")
(require "core.scm")
(require "constants.scm")
(require "numbers.scm")

;;; This module provides a procedure for generating code from a parse
;;; tree. The compilation of a Scheme module proceeds by compiling each
;;; lambda into a private static method containing the code that
;;; implements that method. An instance of the Java module is also
;;; stored as a public static field, which pairs the procedure with an
;;; environment. Each procedure is given a unique index value that is
;;; used in the module's apply method to look it up.

(provide cg)

;; TODO: We will eventually need to maitain a symbol table while 
;; compiling. This is especially needed in places like the set! special
;; form where the variable being set must already be defined. This is
;; something that should be caught at compile time. In addition we will
;; be able to check for unbound variable errors at compile time, which 
;; is definitely good.

;; TODO: I think that the decision will be to still use heap-allocated
;; frames for the procedure environments at run-time, but to also still
;; add a compile-time environment (i.e., symbol table) so that some 
;; errors can be checked at compile time (e.g., arity, set!...).

(define *init* '())
;; Each module is given a private static init method that is used to 
;; create and initialize Scheme procedures as they are compiled.

(define *procedures* '())
;; Java names of each compiled procedure stored in *procedures*.

(define *cf* '())
;; Class File object for the module being compiled. 

;;; In general, the following procedures request locals, code, env-src,
;;; and env arguments, even if they might not use them directly. Locals
;;; indicates the first available local variable index for the cg call,
;;; and is used to ensure that nested cg calls don't overwrite another
;;; call's local variable. Code should be the code attribute object to
;;; write any output to. Env-src should indicate which environment to
;;; use for code generated within the procedure If env-src is null
;;; (i.e. (quote ())), the global environment is used. Otherwise,
;;; env-src should be the string name of a static procedure field to
;;; query for an environment. For example, if env-src is "p2", then an
;;; environment would be pushed onto the stack via p2.getEnvironment().
;;; Env is an actual environment object mapping variables => values.

(define (cg cf ast)
  (define (cg-ast ast)
    (or (null? ast)
        (begin (cg-datum (car ast) *init* '() 2)
               (cg-ast (cdr ast)))))
  (let ((access (bitwise-ior ACC-PRIVATE ACC-STATIC)))
    (set! *init* ((cf 'add-method) "init" "()V" access)))
  (set! *procedures* '())
  (set! *cf* cf)
  (cg-ast ast)
  (return *init*)
  (cg-apply-method))

(define (cg-datum datum code env-src locals)
  (cond ((symbol-datum? datum)
         (cg-symbol datum code env-src locals))
        ((boolean-datum? datum)
         (cg-boolean datum code env-src locals))
        ((number-datum? datum)
         (cg-number datum code env-src locals))
        ((character-datum? datum)
         (cg-character datum code env-src locals))
        ((string-datum? datum)
         (cg-string datum code env-src locals))
        ((vector-datum? datum)
         (cg-vector datum code env-src locals))
        ((any-list-datum? datum)
         (cg-list datum code env-src locals))
        (else 
         (error "unrecognized datum -- CG-DATUM" datum))))

(define (cg-datum-quoted datum code env-src locals)
  (cond ((symbol-datum? datum)
         (cg-symbol-quoted datum code env-src locals))
        ((any-list-datum? datum)
         (cg-list-quoted datum code env-src locals))
        ((pair-datum? datum)
         (cg-pair-quoted datum code env-src locals))
        (else
         (cg-datum datum code env-src locals))))

(define (cg-symbol datum code env-src locals)
  (let ((type (format "(Ljava/lang/String;)~a" *obj-type*)))
    (cg-env code env-src)
    (ldc-string (value-datum datum) code)
    (invokevirtual *env-class* "get" type code)))

(define (cg-symbol-quoted datum code env-src locals)
  (let ((type (format "(Ljava/lang/String;)~a" *symbol-type*)))
    (aload 0 code)
    (ldc-string (value-datum datum) code)
    (invokevirtual (*cf* 'name) "addSymbol" type code)))

(define (cg-boolean datum code env-src locals)
  (let ((name (if (equal? (value-datum datum) "#t") "TRUE" "FALSE")))
    (getstatic *bool-class* name *bool-type* code)))

(define (cg-number datum code env-src locals)
  (let ((in (open-input-string (value-datum datum))))
    (cg-number-prefix in code env-src locals)))

(define (cg-number-prefix in code env-src locals)
  (define (radix-or-exact)
    (case (read-char in)
      ((#\b) (radix-1 2))
      ((#\d) (radix-1 10))
      ((#\o) (radix-1 8))
      ((#\x) (radix-1 16))
      ((#\i) (exact-1 #f))
      ((#\e) (exact-1 #t))))
  (define (radix-1 radix)
    (let ((next-char (peek-char in)))
      (if (char=? #\# next-char)
          (begin (read-char in) (exact-2 radix))
          (cg-number-complex (cons '() radix) in code env-src locals))))
  (define (exact-1 exact)
    (let ((next-char (peek-char in)))
      (if (char=? #\# next-char)
          (begin (read-char in) (radix-2 exact))
          (cg-number-complex (cons exact 10) in code env-src locals))))
  (define (radix-2 exact)
    (let ((radix (case (read-char in)
                   ((#\b) 2)
                   ((#\d) 10)
                   ((#\o) 8)
                   ((#\x) 16))))
      (cg-number-complex (cons exact radix) in code env-src locals)))
  (define (exact-2 radix)
    (let ((exact (char=? #\e (read-char in))))
      (cg-number-complex (cons exact radix) in code env-src locals)))
  (let ((next-char (peek-char in)))
    (if (char=? #\# next-char)
        (begin (read-char in)
               (radix-or-exact))
        (cg-number-complex (cons '() 10) in code env-src locals))))

(define (cg-number-complex prefix in code env-src locals)
  (if (cg-number-real prefix in code env-src locals)
      (let ((next-char (peek-char in)))
        (cond ((eof-object? next-char)
               #t)
              ((char-sign? next-char)
               (let ((type (format "(~a~a)V" *real-type* *real-type*)))
                 (astore locals code)
                 (new-class *complex-class* code)
                 (dup code)
                 (aload locals code)
                 (cg-number-real prefix in code env-src locals)
                 (invokespecial *complex-class* "<init>" type code)))
              ((char=? #\@ next-char) #f)
              ;; TODO this is (I think?) the last part of the numerics
              (else
               #f)))
      (let ((next-char (peek-char in))
            (type (format "(~a~a)V" *real-type* *real-type*)))
        (cond ((eof-object? next-char)
               #t)
              ((char=? #\i next-char)
               (astore locals code)
               (new-class *complex-class* code)
               (dup code)
               (getstatic *int-class* "ZERO" *int-type* code)
               (aload locals code)
               (invokespecial *complex-class* "<init>" type code))))))

(define (cg-number-real prefix in code env-src locals)
  (if (char-sign? (peek-char in))
      (cg-number-ureal prefix (read-char in) in code env-src locals)
      (cg-number-ureal prefix #\+ in code env-src locals)))

(define (cg-number-ureal prefix sign in code env-src locals)
  ;; TODO can we break this up at all?
  (let* ((digits (if (char=? #\+ sign)
                     (uinteger prefix in)
                     (format "-~a" (uinteger prefix in))))
         (next-char (peek-char in)))
    (cond ((or (eof-object? next-char) (char=? #\/ next-char))
           (let ((int-init-type "(Ljava/lang/String;IZ)V")
                 (real-init-type
                  (format "(~a~a)V" *int-type* *int-type*)))
             (unless (eof-object? next-char)
                     (new-class *real-class* code)
                     (dup code))
             (new-class *int-class* code)
             (dup code)
             (ldc-string digits code)
             (iconst (cdr prefix) code)
             (let ((exact (if (or (null? (car prefix)) (car prefix)) 1 0)))
               (iconst exact code))
             (invokespecial *int-class* "<init>" int-init-type code)
             (or (eof-object? next-char)
                 (begin 
                   (read-char in)
                   (new-class *int-class* code)
                   (dup code)
                   (ldc-string (uinteger prefix in) code)
                   (iconst (cdr prefix) code)
                   (iconst 1 code)
                   (invokespecial *int-class* "<init>" int-init-type code)
                   (invokespecial *real-class* "<init>" real-init-type code)
                   #t))))
          ((char=? #\. next-char)
           (read-char in)
           (let ((number (if (null? digits) "0" digits)))
             (cg-number-decimal prefix number in code env-src locals)))
          ((char=? #\i next-char)
           (if (zero? (string-length digits))
               (begin (getstatic *int-class* "ONE" *int-type* code)
                      (if (char=? #\- sign)
                          (let ((neg-type (format "()~a" *num-type*)))
                            (dup code)
                            (invokevirtual *int-class* "negate" neg-type code))
                          'ok)
                      #f)
               (let ((init-type "(Ljava/lang/String;I)V"))
                 (new-class *int-class* code)
                 (dup code)
                 (ldc-string digits code)
                 (iconst (cdr prefix) code)
                 (invokespecial *int-class* "<init>" init-type code)
                 (if (char=? #\- sign)
                     (let ((neg-type (format "()~a" *num-type*)))
                       (dup code)
                       (invokevirtual *int-class* "negate" neg-type code))
                     'ok)
                 #f)))
          (else
           (new-class *real-class* code)
           (dup code)
           (ldc-string digits code)
           (iconst (if (or (null? (car prefix)) (car prefix)) 1 0) code)
           (invokespecial *real-class* "<init>" "(Ljava/lang/String;Z)V" code)
           #t))))

(define (cg-number-decimal prefix number in code env-src locals)
  (define (suffix)
    (let ((next-char (peek-char in)))
      (cond ((eof-object? next-char) "")
            ((char-exponent-marker? next-char)
             (read-char in) ;; exponent marker
             (if (char-sign? (peek-char in))
                 (let ((sign (read-char in)))
                   (format "e~a~a" sign (uinteger prefix in)))
                 (format "e~a" (uinteger prefix in))))
            (else ""))))
  (let ((value (format "~a.~a~a" number (uinteger prefix in) (suffix)))
        (exact (if (or (null? (car prefix)) (not (car prefix))) 0 1)))
    (new-class *real-class* code)
    (dup code)
    (ldc-string value code)
    (iconst exact code)
    (invokespecial *real-class* "<init>" "(Ljava/lang/String;Z)V" code)
    #t))

(define (uinteger prefix in)
  (define digits? (radix->digits? (cdr prefix)))
  (define (collect digits)
    (let ((next-char (peek-char in)))
      (if (or (eof-object? next-char)
              (not (digits? next-char)))
          (reverse digits)
          (collect (cons (read-char in) digits)))))
  (list->string (collect '())))

(define (cg-character datum code env-src locals)
  ;; TODO
  #f)

(define (cg-string datum code env-src locals)
  (new-class *str-class* code)
  (dup code)
  (ldc-string (value-datum datum) code)
  (invokespecial *str-class* "<init>" "(Ljava/lang/String;)V" code)) 

(define (cg-vector datum code env-src locals)
  (new-class *vector-class* code)
  (dup code)
  (iconst (length (value-datum datum)) code)
  (anewarray-class *obj-class* code)
  (let ((type (format "([~a)V" *obj-type*))
        (i 0))
    (for-each
     (lambda (el)
       (dup code)
       (iconst i code)
       (cg-datum-quoted el code env-src locals)
       (aastore code)
       (set! i (+ i 1)))
     (value-datum datum))
    (invokespecial *vector-class* "<init>" type code)))

(define (cg-pair-quoted datum code env-src locals)
  (cg-pair-<quoted> datum code env-src locals cg-datum-quoted))

(define (cg-pair-<quoted> datum code env-src locals <cg-quoted>)
  (let ((pair (value-datum datum))
        (type (format "(~a~a)V" *obj-type* *obj-type*)))
    (new-class *pair-class* code)
    (dup code)
    (<cg-quoted> (car pair) code env-src locals)
    (<cg-quoted> (cdr pair) code env-src locals)
    (invokespecial *pair-class* "<init>" type code)))  

(define (cg-list datum code env-src locals)
  (cond ((begin-datum? datum)
         (cg-begin datum code env-src locals))
        ((define-datum? datum)
         (cg-define datum code env-src locals))
        ((if-datum? datum)
         (cg-if datum code env-src locals))
        ((lambda-datum? datum)
         (cg-lambda datum code env-src locals))
        ((quote-datum? datum)
         (cg-quote datum code env-src locals))
        ((set!-datum? datum)
         (cg-set! datum code env-src locals))
        (else
         (cg-application datum code env-src locals))))

(define (cg-list-quoted datum code env-src locals)
  (cg-list-<quoted> datum code env-src locals cg-datum-quoted))

(define (cg-list-<quoted> datum code env-src locals <cg-quoted>)
  (let ((seq (value-datum datum)))
    (if (null? seq)
        (getstatic *pair-class* "EMPTY_LIST" *pair-type* code)
        (let ((pair-init-type (format "(~a~a)V" *obj-type* *obj-type*))
              (set-cdr-type (format "(~a)V" *obj-type*))
              (it (+ 1 locals))
              (el (+ 2 locals))
              (next-locals (+ 3 locals)))
          (new-class *pair-class* code)
          (dup code)
          (<cg-quoted> (car seq) code env-src next-locals)
          (aconst_null code)
          (invokespecial *pair-class* "<init>" pair-init-type code)
          (astore locals code)
          (aload locals code)
          (astore it code)
          (do ((seq (cdr seq) (cdr seq))
               (i (- (length seq) 2) (- i 1)))
              ((null? seq) #t)
            (aload it code)
            (if (and (pair-datum? (car seq))
                     (zero? i)
                     (dotted-list-datum? datum))
                (<cg-quoted> (car seq) code env-src next-locals)
                (begin 
                  (new-class *pair-class* code)
                  (dup code)
                  (<cg-quoted> (car seq) code env-src next-locals)
                  (aconst_null code)
                  (invokespecial *pair-class* "<init>" pair-init-type code)))
            (astore el code)
            (aload el code)
            (invokevirtual *pair-class* "setCdr" set-cdr-type code)
            (aload el code)
            (astore it code))
          (if (dotted-list-datum? datum)
              (aload locals code)
              (begin
                (aload it code)
                (getstatic *pair-class* "EMPTY_LIST" *pair-type* code)
                (invokevirtual *pair-class* "setCdr" set-cdr-type code)
                (aload locals code)))))))

(define (cg-begin datum code env-src locals)
  (cg-begin-seq (cdr (value-datum datum)) code env-src locals))

(define (cg-begin-seq seq code env-src locals)
  (if (null? (cdr seq))
      (cg-datum (car seq) code env-src locals)
      (begin (cg-datum (car seq) code env-src locals)
             (pop code)
             (cg-begin-seq (cdr seq) code env-src locals))))

(define (cg-define datum code env-src locals)
  (cg-env code env-src)
  (ldc-string (value-datum (define-variable datum)) code)
  (cg-datum (define-value datum) code env-src locals)
  (cg-env-put code)
  (if (define-main-datum? datum env-src)
      (cg-define-main datum code env-src locals)
      #t))

(define (define-main-datum? datum env-src)
  (and (equal? "main" (value-datum (define-variable datum)))
       (lambda-datum? (define-value datum))
       (null? env-src)))

(define (cg-define-main datum code env-src locals)
  (let* ((access (bitwise-ior ACC-PUBLIC ACC-STATIC))
         (main ((*cf* 'add-method) "main" "([Ljava/lang/String;)V" access)))
    (cg-env main env-src)
    (ldc-string "main" main)
    (cg-env-get main)
    (checkcast *proc-class* main)
    (iconst 1 main)
    (anewarray-class *obj-class* main)
    (dup main)
    (iconst 0 main)
    (aload 0 main)
    (let ((type (format "([Ljava/lang/String;)~a" *obj-type*)))
      (invokestatic *module-class* "toList" type main))
    (aastore main)
    (let ((type (format "([~a)~a" *obj-type* *obj-type*)))
      (invokevirtual *proc-class* "apply" type main)
      (return main))))

(define (cg-if datum code env-src locals)
  (getstatic *bool-class* "FALSE" *bool-type* code)
  (cg-datum (if-test datum) code env-src locals)
  (let ((code-output ((code 'code-output)))
        (alt-output (open-output-bytes))
        (con-output (open-output-bytes)))
    ((code 'set-code-output!) alt-output)
    (cg-datum (if-alternate datum) code env-src locals)
    ((code 'set-code-output!) con-output)
    (cg-datum (if-consequent datum) code env-src locals)
    ((code 'set-code-output!) code-output)
    (let ((alternate (get-output-bytes alt-output #t))
          (consequent (get-output-bytes con-output #t)))
      ;; 3 bytes for if_acmpne + 3 bytes for the extra goto = 6 bytes offset
      (if_acmpne (+ 6 (bytes-length alternate)) code)
      (write-bytes alternate code-output)
      (goto (+ 3 (bytes-length consequent)) code)
      (write-bytes consequent code-output))))

(define (cg-lambda datum code env-src locals)
  (define (cg-arg i arg proc proc-name var?)
    (let ((p (if (pair-datum? arg) (car-datum arg) arg)))
      (cg-env proc proc-name)
      (ldc-string (value-datum p) proc)
      (aload 1 proc)
      (iconst i proc)
      (if var?
          (let ((type (format "([~aI)~a" *obj-type* *obj-type*)))
            (invokestatic *module-class* "toList" type proc))
          (aaload proc))
      (cg-env-sput proc)
      (pop proc)
      (if (pair-datum? arg)
          (cg-arg (+ i 1) (cdr-datum arg) proc proc-name #t)
          (+ i 1))))
  (let ((proc-name (index->name (length *procedures*)))
        (proc (make-procedure (lambda-arity datum) env-src))
        (type (format "(Ljava/lang/String;~a)~a" *obj-type* *obj-type*))
        (env-is-type (format "(~a)V" *env-type*))
        (i 0))
    (aload 0 proc)
    (invokevirtual (*cf* 'name) "enter" "()V" proc)
    (for-each ;; put args in environment
     (lambda (p) (set! i (cg-arg i p proc proc-name #f)))
     (let ((params (lambda-parameters datum)))
       (if (symbol-datum? params)
           (list params)
           (value-datum params))))
    (cg-begin-seq (lambda-body datum) proc proc-name locals)
    ;; TODO: Check for TAILCALL and if needed, JUMP
    (aload 0 proc)
    (invokevirtual (*cf* 'name) "exit" "()V" proc)
    (areturn proc)
    (getstatic (*cf* 'name) proc-name (class->type (*cf* 'name)) code)
    (if (not (null? env-src))
        (let ((get-env-type (format "()~a" *env-type*)))
          (dup code)
          (aload 0 code)
          (invokevirtual (*cf* 'name) "getEnvironment" get-env-type code)
          (invokevirtual (*cf* 'name) "setEnvironment" env-is-type code))
        'ok)))

(define (index->name index)
  (format "p~a" index))

(define (make-procedure arity env-src)
  ;; arity should be (cons min max) where max may be 'MAX-VALUE
  (let ((proc-name (index->name (length *procedures*)))
        (proc-index (length *procedures*)))
    (set! *procedures* (append *procedures* (list proc-name)))
    (make-procedure-field! proc-name proc-index arity env-src)
    (make-procedure-method proc-name proc-index arity env-src)))

(define (make-procedure-field! proc-name proc-index arity env-src)
  (let ((proc-class (*cf* 'name))
        (proc-type (class->type (*cf* 'name)))
        (access (bitwise-ior ACC-PUBLIC ACC-STATIC ACC-FINAL)))
    (new-class proc-class *init*)
    (dup *init*)
    (iconst (car arity) *init*)
    (if (eq? 'MAX-VALUE (cdr arity))
        (getstatic "java/lang/Integer" "MAX_VALUE" "I" *init*)
        (iconst (cdr arity) *init*))
    (iconst proc-index *init*)
    (cg-env *init* env-src)
    (let ((init-type (format "(III~a)V" *env-type*)))
      (invokespecial proc-class "<init>" init-type *init*))
    (putstatic proc-class proc-name proc-type *init*)
    ((*cf* 'add-field!) proc-name proc-type access)))

(define (make-procedure-method proc-name proc-index arity env-src)
  (let ((proc-type (class->type (*cf* 'name)))
        (access (bitwise-ior ACC-PUBLIC ACC-STATIC)))
    (let* ((type (format "(~a[~a)~a" proc-type *obj-type* *obj-type*))
           (code ((*cf* 'add-method) proc-name type access)))
      (aload 0 code)
      (aload 1 code)
      (arraylength code)
      (invokevirtual (*cf* 'name) "checkArity" "(I)V" code)
      code)))

(define (cg-quote datum code env-src locals)
  (cg-datum-quoted (quote-datum datum) code env-src locals))

(define (cg-set! datum code env-src locals)
  ;; TODO: check that the variable has already been defined
  (cg-env code env-src)
  (ldc-string (value-datum (set!-variable datum)) code)
  (cg-datum (set!-value datum) code env-src locals)
  (cg-env-put code))

(define (cg-application datum code env-src locals)
  ;; TODO: make loop in bytecode to build args array instead of duplicating 
  (let ((proc (car (value-datum datum)))
        (type (format "([~a)~a" *obj-type* *obj-type*)))
    (cg-datum proc code env-src locals)
    (checkcast *proc-class* code)
    (let ((args (cdr (value-datum datum)))
          (i 0))
      (iconst (length args) code)
      (anewarray-class *obj-class* code)
      (for-each
       (lambda (arg)
         (dup code)          ;; dup because we need to leave the
         (iconst i code)     ;; args array on top of the stack
         (cg-datum arg code env-src locals)
         (aastore code)
         (set! i (+ i 1)))
       args))
    (invokevirtual *proc-class* "apply" type code)))

(define (cg-env code env-src)
  (if (null? env-src)
      (cg-global-env code)
      (let ((type (class->type (*cf* 'name)))
            (get-type (format "()~a" *env-type*)))
        (getstatic (*cf* 'name) env-src type code)
        (invokevirtual (*cf* 'name) "getEnvironment" get-type code))))

(define (cg-global-env code)
  (getstatic (*cf* 'name) "GLOBAL_ENV" *env-type* code))

(define (cg-env-get code)
  (let ((type (format "(Ljava/lang/String;)~a" *obj-type*)))
    (invokevirtual *env-class* "get" type code)))

(define (cg-env-put code)
  (generic-cg-env-put code "put"))

(define (cg-env-sput code)
  (generic-cg-env-put code "sput"))

(define (generic-cg-env-put code method)
  (let ((type (format "(Ljava/lang/String;~a)~a" *obj-type* *obj-type*)))
    (invokevirtual *env-class* method type code)))

(define (cg-apply-method)
  (define (make-offset-list base width)
    (let ((offset base))
      (map (lambda (p)
             (let ((o offset))
               (set! offset (+ offset width))
               o))
           *procedures*)))
  ;; tableswitch begins at address 4 => 3 null padding bytes are used
  ;; 1 (opcode) + 3 (padding) + 12 (default,low,high) = 16 bytes before offsets
  ;; (* 4 (length *procedures*)) bytes used for offsets
  ;; .: 16 + offsets = address of first case in switch statement
  (let* ((type (format "([~a)~a" *obj-type* *obj-type*))
         (code ((*cf* 'add-method) "apply" type ACC-PUBLIC))
         (base (+ 16 (* 4 (length *procedures*)))))
    (let ((offsets (make-offset-list base 6))
          (default (+ base (* 6 (length *procedures*)))))
      (aload 0 code)                                     ; address 0
      (invokevirtual (*cf* 'name) "getIndex" "()I" code) ; address 1,2,3
      (tableswitch 3 default 0 (- (length *procedures*) 1) offsets code))
    (let* ((this-type (class->type (*cf* 'name)))
           (type (format "(~a[~a)~a" this-type *obj-type* *obj-type*)))
      (for-each
       (lambda (p)
         (aload 0 code) 
         (aload 1 code) 
         (invokestatic (*cf* 'name) p type code)
         (areturn code))
       (append *procedures* (list "cd"))))))
